import os
import time
from datetime import datetime, timedelta
import pandas as pd
import numpy as np

DAILY_CSV = 'stock_daily_data.csv'
STOCKS_CSV = 'stocks.csv'
TECH_CSV = 'stock_technical_indicators.csv'
TECH_CLEAN_CSV = 'stock_technical_indicators_clean.csv'
PRED_CSV = 'ai_predictions.csv'
STOCKS_FILE = 'stocks_list.txt'

STOCK_BASIC_FIELDS = ['ts_code','symbol','name','area','industry','list_date']

def read_daily():
    if not os.path.exists(DAILY_CSV):
        raise FileNotFoundError(f'未找到日线数据文件: {DAILY_CSV}')
    df = pd.read_csv(DAILY_CSV)
    # 统一列类型
    for col in ['open','high','low','close','vol','amount']:
        if col in df.columns:
            df[col] = pd.to_numeric(df[col], errors='coerce')
    df['trade_date'] = df['trade_date'].astype(str)
    # 按 ts_code+trade_date 去重保留最后
    if 'id' in df.columns:
        df = df.drop(columns=['id'])
    df = df.sort_values(['ts_code','trade_date']).drop_duplicates(['ts_code','trade_date'], keep='last')
    df = df.reset_index(drop=True)
    df['id'] = range(1, len(df)+1)
    df = df[['id','ts_code','trade_date','open','high','low','close','vol','amount']]
    df.to_csv(DAILY_CSV, index=False)
    return df

def compute_indicators(daily_df: pd.DataFrame):
    # 以最近 ~60 天为窗口计算
    cutoff = (datetime.now() - timedelta(days=60)).strftime('%Y%m%d')
    d = daily_df.copy()
    d['trade_date'] = d['trade_date'].astype(str)
    d = d[d['trade_date'] >= cutoff]
    rows = []
    for code, grp in d.groupby('ts_code'):
        g = grp.sort_values('trade_date')
        close_s = pd.to_numeric(g['close'], errors='coerce')
        ma5 = close_s.rolling(window=5, min_periods=5).mean()
        ma20 = close_s.rolling(window=20, min_periods=20).mean()
        delta = close_s.diff()
        gain = delta.clip(lower=0)
        loss = -delta.clip(upper=0)
        avg_gain = gain.ewm(alpha=1/14, min_periods=14, adjust=False).mean()
        avg_loss = loss.ewm(alpha=1/14, min_periods=14, adjust=False).mean()
        rs = avg_gain / (avg_loss.replace(0, pd.NA))
        rsi = 100 - (100 / (1 + rs))
        ema12 = close_s.ewm(span=12, adjust=False).mean()
        ema26 = close_s.ewm(span=26, adjust=False).mean()
        macd = ema12 - ema26
        mid = close_s.rolling(window=20, min_periods=20).mean()
        std = close_s.rolling(window=20, min_periods=20).std()
        boll_upper = mid + 2 * std
        boll_lower = mid - 2 * std
        out = pd.DataFrame({
            'ts_code': g['ts_code'],
            'trade_date': g['trade_date'],
            'ma5': ma5.round(2),
            'ma20': ma20.round(2),
            'rsi': rsi.round(2),
            'macd': macd.round(4),
            'boll_upper': boll_upper.round(2),
            'boll_lower': boll_lower.round(2)
        })
        rows.append(out)
    if rows:
        tech = pd.concat(rows, ignore_index=True)
        tech = tech.sort_values(['ts_code','trade_date']).reset_index(drop=True)
        tech['id'] = range(1, len(tech)+1)
        tech = tech[['id','ts_code','trade_date','ma5','ma20','rsi','macd','boll_upper','boll_lower']]
    else:
        tech = pd.DataFrame(columns=['id','ts_code','trade_date','ma5','ma20','rsi','macd','boll_upper','boll_lower'])
    tech.to_csv(TECH_CSV, index=False)
    clean = tech.dropna(subset=['ma5','ma20','rsi','macd','boll_upper','boll_lower'], how='any')
    clean.to_csv(TECH_CLEAN_CSV, index=False)
    return tech, clean

def build_stocks_from_daily(daily_df: pd.DataFrame):
    # 从 daily 中取 ts_code，生成最简 stocks 表；name 尝试用 AkShare 映射
    try:
        import akshare as ak
    except Exception:
        ak = None
    ts_codes = sorted(daily_df['ts_code'].dropna().unique().tolist())
    name_map = {}
    if ak is not None:
        # 尝试多次获取映射
        for i in range(3):
            try:
                raw = ak.stock_info_a_code_name()
                cols = list(raw.columns)
                code_candidates = ['symbol','股票代码','代码','code']
                name_candidates = ['name','股票简称','简称','security_name']
                sym_col = next((c for c in code_candidates if c in cols), None)
                name_col = next((c for c in name_candidates if c in cols), None)
                if sym_col and name_col:
                    raw[sym_col] = raw[sym_col].astype(str).str.strip()
                    raw[name_col] = raw[name_col].astype(str).str.strip()
                    name_map = dict(zip(raw[sym_col], raw[name_col]))
                    break
            except Exception:
                time.sleep(2)
    rows = []
    for c in ts_codes:
        sym = c.split('.')[0]
        nm = name_map.get(sym)
        rows.append({'ts_code': c, 'symbol': sym, 'name': nm if nm else None, 'area': None, 'industry': None, 'list_date': None})
    df = pd.DataFrame(rows, columns=STOCK_BASIC_FIELDS)
    df.insert(0, 'id', range(1, len(df)+1))
    df.to_csv(STOCKS_CSV, index=False)
    return df

def build_predictions_from_indicators(tech_df: pd.DataFrame):
    # 使用技术指标生成明日预测分数，覆盖式重建
    predict_date = datetime.now().strftime('%Y%m%d')
    for_date = (datetime.now() + timedelta(days=1)).strftime('%Y%m%d')
    cols_needed = {'ts_code','trade_date','ma5','ma20','rsi'}
    tech = tech_df.copy()
    missing = cols_needed - set(tech.columns)
    if missing:
        # 指标不齐时返回空表
        out = pd.DataFrame(columns=['id','ts_code','predict_date','for_date','prediction_score'])
        out.to_csv(PRED_CSV, index=False)
        return out
    tech = tech.dropna(subset=['ma5','ma20','rsi'], how='any')
    if tech.empty:
        out = pd.DataFrame(columns=['id','ts_code','predict_date','for_date','prediction_score'])
        out.to_csv(PRED_CSV, index=False)
        return out
    tech['trade_date'] = tech['trade_date'].astype(str)
    # 取每只股票最近一日指标
    tech_latest = tech.sort_values(['ts_code','trade_date']).groupby('ts_code', as_index=False).tail(1)
    # MA 差信号：tanh 压缩 + 映射到 [0,1]
    ma_diff = (pd.to_numeric(tech_latest['ma5'], errors='coerce') - pd.to_numeric(tech_latest['ma20'], errors='coerce')) / pd.to_numeric(tech_latest['ma20'], errors='coerce')
    ma_diff = ma_diff.replace([np.inf, -np.inf], np.nan).fillna(0.0)
    score_ma = 0.5 * (np.tanh(ma_diff * 5.0) + 1.0)
    # RSI 信号：以 50 为中枢线性归一
    rsi = pd.to_numeric(tech_latest['rsi'], errors='coerce').fillna(50.0)
    score_rsi = np.clip(0.5 + (rsi - 50.0) / 100.0, 0.0, 1.0)
    # 综合得分：等权平均
    score = np.round(0.5 * score_ma + 0.5 * score_rsi, 4)
    preds = pd.DataFrame({
        'ts_code': tech_latest['ts_code'].astype(str),
        'predict_date': predict_date,
        'for_date': for_date,
        'prediction_score': score
    })
    preds = preds.reset_index(drop=True)
    preds.insert(0, 'id', range(1, len(preds)+1))
    preds = preds[['id','ts_code','predict_date','for_date','prediction_score']]
    preds.to_csv(PRED_CSV, index=False)
    return preds

def main():
    daily = read_daily()
    tech, clean = compute_indicators(daily)
    stocks = build_stocks_from_daily(daily)
    preds = build_predictions_from_indicators(tech)
    print('[DONE] 五张表已根据现有日线数据完成重建:')
    print(STOCKS_CSV, DAILY_CSV, TECH_CSV, TECH_CLEAN_CSV, PRED_CSV)

if __name__ == '__main__':
    main()
